
module Kalman
        ( KalmanState
        , kfUpdate
        , kfInitial1D
        , findKalmanEstimate
        )
        where

data KalmanState a = KalmanState {
        ksXhat      :: a,
        ksP         :: a,
        ksK         :: a,
        ksR         :: a,
        ksQ         :: a
        } deriving (Show)

kfInitial1D ::  KalmanState Double
kfInitial1D = KalmanState { ksXhat = 0.0, ksP = 1.0, ksK = 0.0, ksR = 0.1, ksQ = 1e-5 }

class Fractional a=>KalmanFilter a where
        kfUpdate :: KalmanState a->a->KalmanState a
        kfUpdate f z = KalmanState xhat p k (ksR f) (ksQ f)
                where   pminus  = (ksP f) + (ksQ f)
                        k       = pminus/(pminus + (ksR f))
                        xhat    = (ksXhat f) + k * (z - (ksXhat f))
                        p       = (1-k)*pminus


instance KalmanFilter Double

findKalmanEstimate ::  [(t, [Double])] -> [KalmanState Double]
findKalmanEstimate x = foldl threeUpdates threeKf x
        where   threeKf = replicate 3 kfInitial1D
                threeUpdates [k1, k2, k3] (_, [z1, z2, z3]) = [kfUpdate k1 z1, kfUpdate k2 z2, kfUpdate k3 z3]

